
import argparse
import logging
import os
import sys
import numpy as np
import torch
import torch.nn as nn
from torch import optim
from tqdm import tqdm
from eval import eval_net
from unet3d.model import UNet2D, UNet3D
import matplotlib.pyplot as plt 
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader, random_split
from dataset import BasicDataset

data_dir = "./DSB_data/"
dir_checkpoint = 'checkpoints/'
epochs = 50
batch_size = 2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
lr = 0.0001
# 验证集百分比
val_percent = 0.1
img_size = (256, 256)
final_sigmoid = True 
in_channels = 1
out_channels = 1

def train_net(net):
    # 分割数据集～ 训练集和验证集
    dataset = BasicDataset(data_dir, img_size=img_size)
    n_val = int(len(dataset) * val_percent)
    n_train = len(dataset) - n_val
    train, val = random_split(dataset, [n_train, n_val])

    train_loader = DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)
    val_loader = DataLoader(val, batch_size=1, shuffle=False, num_workers=8, pin_memory=True, drop_last=True)

    writer = SummaryWriter("./log/", comment=f'LR_{lr}_BS_{batch_size}')
    global_step = 0

    logging.info(f'''Starting training:
        Epochs:          {epochs}
        Batch size:      {batch_size}
        Learning rate:   {lr}
        Training size:   {n_train}
        Validation size: {n_val}
        Device:          {device.type}
        Images size:  {img_size}
    ''')

    optimizer = optim.Adam(net.parameters(), lr=lr, weight_decay=1e-5)
    # optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.9)
    # scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min' if net.n_classes > 1 else 'max', patience=2)
    if final_sigmoid :
        # 证明是二分类
        criterion = nn.BCEWithLogitsLoss()
    else:
        # 多分类
        criterion = nn.CrossEntropyLoss()
    # report_loss = 0
    for epoch in range(epochs):
        net.train()

        epoch_loss = 0
        with tqdm(total=n_train, desc=f'Epoch {epoch + 1}/{epochs}', unit='img') as pbar:
            for batch in train_loader:
                imgs = batch['image']
                true_masks = batch['mask']
               
                # imgs = torch.squeeze(imgs, 2)
                # plt.imshow(imgs)
                # plt.show()
                assert imgs.shape[1] == in_channels, \
                    f'Network has been defined with {in_channels} input channels, ' \
                    f'but loaded images have {imgs.shape[1]} channels. Please check that ' \
                    'the images are loaded correctly.'

                imgs = imgs.to(device=device, dtype=torch.float32)
                mask_type = torch.float32 if out_channels == 1 else torch.long
                true_masks = true_masks.to(device=device, dtype=mask_type)

                masks_pred = net(imgs)
                
                loss = criterion(masks_pred, true_masks)
                epoch_loss += loss.item()
                writer.add_scalar('Loss/train', loss.item(), global_step)

                pbar.set_postfix(**{'loss (batch)': loss.item()})

                optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_value_(net.parameters(), 0.1)
                optimizer.step()
                # report_loss += loss.item()
                pbar.update(imgs.shape[0])
                global_step += 1
                
                if global_step % (n_train // (5 * batch_size)) == 0:
                    # print("report loss is : " + str(report_loss))
                    # report_loss = 0
                    for tag, value in net.named_parameters():
                        tag = tag.replace('.', '/')
                        writer.add_histogram('weights/' + tag, value.data.cpu().numpy(), global_step)
                        writer.add_histogram('grads/' + tag, value.grad.data.cpu().numpy(), global_step)
                    val_score = eval_net(net, val_loader, device)
                    print(val_score)
                    # scheduler.step(val_score)
                    # writer.add_scalar('learning_rate', optimizer.param_groups[0]['lr'], global_step)

                    if out_channels > 1:
                        logging.info('Validation cross entropy: {}'.format(val_score))
                        writer.add_scalar('Loss/test', val_score, global_step)
                    else:
                        logging.info('Validation Dice Coeff: {}'.format(val_score))
                        writer.add_scalar('Dice/test', val_score, global_step)
                    
                    # 三维变二维图像
                    imgs = imgs.squeeze(2)
                    writer.add_images('images', imgs, global_step)
                    if out_channels == 1:
                        true_masks = true_masks.squeeze(2)
                        masks_pred = masks_pred.squeeze(2)
                        writer.add_images('masks/true', true_masks, global_step)
                        writer.add_images('masks/pred', torch.sigmoid(masks_pred) > 0.5, global_step)

        print("epoch loss is : " + str(epoch_loss))
        epoch_loss = 0
        try:
            os.mkdir(dir_checkpoint)
            logging.info('Created checkpoint directory')
        except OSError:
            pass
        torch.save(net.state_dict(),
                    dir_checkpoint + f'unet2D_epoch{epoch + 1}.pth')
        logging.info(f'Checkpoint {epoch + 1} saved !')

    writer.close()


if __name__ == '__main__':
   
    logging.info(f'Using device {device}')

    ## 因为输入是灰度图，所以in channels为1，因为是二分类去分割，因此输出channel 也为 1
    net = UNet2D(in_channels=in_channels, out_channels=out_channels, final_sigmoid=final_sigmoid) # 训练阶段 testing = False 最后一层不会使用sigmoid进行激活。
   
    net.to(device=device)
    # faster convolutions, but more memory
    # cudnn.benchmark = True

    try:
        train_net(net=net)
    except KeyboardInterrupt:
        torch.save(net.state_dict(), 'INTERRUPTED.pth')
        logging.info('Saved interrupt')
        try:
            sys.exit(0)
        except SystemExit:
            os._exit(0)
